#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


if __name__ == '__main__':
    Client_eLMSG, Collection_eLMSG = InitMongoDB("10.188.188.22", 27017, "eLMSG", "eLMSG_20180212")
    Client_MicrobiomeGene, Collection_MicrobiomeGene = InitMongoDB("10.188.188.22", 27017, "Pre_Data", "MicrobiomeGene")

    for result in Collection_eLMSG.find({"Rank": "Species"}, {"Name": 1}):
        # Collection_eLMSG.update_one({"_id":result["_id"]},{"$unset":{"FunctionAnnotaion":1}})
        if "Name" in result:
            count = Collection_MicrobiomeGene.count({"org_ref_taxname": result["Name"]})
            if count != 0:
                Collection_eLMSG.update_one({"_id": result["_id"]}, {"$set": {"FunctionAnnotation.EntrezGene": count}})
                print(result["Name"] + "=" + str(count))
    CloseMongoDB(Client_eLMSG)
    CloseMongoDB(Client_MicrobiomeGene)
